#!/usr/bin/env python

from zmq_socket import ZmqClientReq,ZmqClientSub
import pdb
from time import sleep
import re

class StateMachineConn(object):
	""""""
	def __init__(self,trigger_port,sm_port,sm_state_port):
		"""Connect to state machine and listen to end-signal.
		trigger_ip&trigger_port is corresponding to GUI.

		"""
		sm_ip = "localhost"
		trigger_ip = "localhost" # TM is usually run with dummymachine2
		self.statemachine_heartbeat_ = ZmqClientSub(sm_ip,sm_port)
		self.statemachine_heartbeat_.setFilter("101") # 101 is the ID of Heartbear.
		self.trigger_conn_ = ZmqClientReq(trigger_ip,trigger_port)
		self.statemachine_sub_ = ZmqClientSub(sm_ip,sm_state_port)
		self.statemachine_sub_.setFilter("")

	def startTM(self):
		# Wait for statemachine online (i.e: heart_beat).
		print("Verify Statemachine Heartbeat...")
		self.statemachine_heartbeat_.recv()
		sleep(2)
		print("Verify Statemachine Heartbeat success, begin test...")
		# Check TM Online
		start_msg = """{"uutinfo" : [{"uut" : 0,"sn" : ""}],"function" : "START"}"""
		#start_msg = "START_OK"
		self.trigger_conn_.send(start_msg)
		return self.trigger_conn_.recv()

	def waitForEnd(self):
		# Wait end flag for statemachine.
		end_flag = "SM_IDLE"

		print("Wait end-flag for StateMahine...")
		while(True):
			sm_state = self.statemachine_sub_.recv()
			if (None!=re.search(end_flag,sm_state,flags=0)):
				break
